import os
import datetime
import win32evtlog # 用于读取Windows事件日志
import win32evtlogutil
import wmi # 用于WMI查询，例如获取系统信息
import struct # 用于解析二进制文件，如Prefetch
import glob # 用于查找文件
import sys # 导入sys模块，用于访问命令行参数

def format_timestamp(timestamp):
    """格式化时间戳为可读字符串"""
    return timestamp.strftime("%Y-%m-%d %H:%M:%S")

def get_event_logs_for_day(target_date):
    """
    获取指定日期的Windows事件日志。
    这个函数直接读取活动的事件日志。
    为了避免卡死，我们增加了一个读取限制和进度提示。
    对于离线或更详细的解析，可能需要python-evtx库。
    """
    print(f"\n--- 正在获取 {target_date.strftime('%Y-%m-%d')} 的事件日志 ---")
    # 常见的日志类型，如果安装了Sysmon，可以添加 "Microsoft-Windows-Sysmon/Operational"
    log_types = ["System", "Application", "Security"]

    events_found = 0
    # 存储找到的事件，以便统一输出，减少循环内的打印开销
    daily_events = []

    for log_type in log_types:
        hand = None # 初始化句柄
        try:
            hand = win32evtlog.OpenEventLog(None, log_type)
            total_events = win32evtlog.GetNumberOfEventLogRecords(hand)
            print(f"  日志类型: {log_type}, 总事件数: {total_events}")

            # 从最新的事件开始读取
            flags = win32evtlog.EVENTLOG_BACKWARDS_READ | win32evtlog.EVENTLOG_SEQUENTIAL_READ
            offset = 0
            read_count = 0
            max_read_events_per_log = 100000 # 限制每个日志类型最多读取的事件数量，防止无限循环
            # 注意：这个限制意味着你可能无法获取到非常旧的、但符合日期的事件

            while True:
                # 读取一批事件 (例如，每次读取1000条)
                events = win32evtlog.ReadEventLog(hand, flags, offset)
                if not events:
                    break

                for event in events:
                    read_count += 1
                    event_time = event.TimeGenerated.replace(tzinfo=None) # 移除时区信息进行比较

                    # 检查事件是否在目标日期内
                    if event_time.date() == target_date.date():
                        source = event.SourceName
                        event_id = event.EventID
                        # 延迟 SafeFormatMessage 调用，只在确认需要时才调用，或在收集后统一格式化
                        # 这里我们仍然在循环中格式化，但可以优化
                        event_string = win32evtlogutil.SafeFormatMessage(event, log_type)
                        daily_events.append(f"    时间: {format_timestamp(event_time)}, 来源: {source}, ID: {event_id}, 描述: {event_string.strip().splitlines()[0]}")
                        events_found += 1
                    elif event_time.date() < target_date.date():
                        # 如果事件时间早于目标日期，说明已经读过目标日期的所有事件，可以停止
                        # 只有在以 BACKWARDS_READ 模式时才有效
                        break

                    if read_count % 5000 == 0: # 每读取5000条事件打印一次进度
                        print(f"    - 已处理 {read_count} 条事件...", end='\r')

                    if read_count >= max_read_events_per_log:
                        print(f"\n    - 达到 {max_read_events_per_log} 条事件读取上限，停止当前日志类型读取。")
                        break
                else: # 如果内层循环没有break，说明这批事件都处理完了，继续读取下一批
                    continue
                break # 如果内层循环break了，说明已经读到目标日期之前的事件或达到上限，跳出外层循环

        except Exception as e:
            print(f"\n  无法读取 {log_type} 日志: {e}")
        finally:
            if hand:
                win32evtlog.CloseEventLog(hand)

    # 统一输出找到的事件
    print(f"\n--- 找到 {events_found} 条 {target_date.strftime('%Y-%m-%d')} 的事件日志 ---")
    for event_line in daily_events:
        print(event_line)


def get_prefetch_files_for_day(target_date):
    """
    解析Prefetch文件，获取应用程序运行信息。
    Prefetch文件通常位于 C:\\Windows\\Prefetch。
    注意：Prefetch文件格式复杂，这里仅展示最简单的最后运行时间解析。
    更详细的解析需要专门的库或对PF文件结构有深入了解。
    """
    print(f"\n--- 正在扫描 {target_date.strftime('%Y-%m-%d')} 的Prefetch文件 ---")
    prefetch_dir = os.path.join(os.environ.get('SystemRoot', 'C:\\Windows'), 'Prefetch')
    if not os.path.exists(prefetch_dir):
        print(f"  Prefetch目录不存在: {prefetch_dir}")
        return

    pf_files_found = 0
    # Prefetch文件通常以 .pf 结尾
    for pf_file in glob.glob(os.path.join(prefetch_dir, "*.pf")):
        try:
            with open(pf_file, 'rb') as f:
                # Prefetch文件头通常包含一些时间戳信息
                # 这是一个非常简化的示例，PF文件结构远比这复杂
                # 假设我们能从某个固定偏移量读取一个NTFS时间戳（FILETIME）
                f.seek(120) # 这是一个示例偏移量，具体取决于PF文件版本和类型
                filetime_bytes = f.read(8)
                if len(filetime_bytes) == 8:
                    filetime = struct.unpack('<Q', filetime_bytes)[0]
                    # 将Windows FILETIME转换为Python datetime对象
                    # FILETIME是自1601年1月1日以来的100纳秒间隔数
                    dt_object = datetime.datetime(1601, 1, 1) + datetime.timedelta(microseconds=filetime / 10)

                    if dt_object.date() == target_date.date():
                        app_name = os.path.basename(pf_file).split('-')[0] # 简单的提取应用名
                        print(f"    应用: {app_name}, 最后运行时间: {format_timestamp(dt_object)}")
                        pf_files_found += 1
        except Exception as e:
            # print(f"  解析Prefetch文件 {os.path.basename(pf_file)} 失败: {e}")
            pass # 忽略解析失败的Prefetch文件
    print(f"--- 找到 {pf_files_found} 个 {target_date.strftime('%Y-%m-%d')} 相关Prefetch文件 ---")


def get_lnk_files_for_day(target_date):
    """
    扫描用户桌面、最近文档等位置的.lnk文件，获取其时间戳。
    LNK文件解析也比较复杂，这里只检查文件本身的修改时间。
    """
    print(f"\n--- 正在扫描 {target_date.strftime('%Y-%m-%d')} 的快捷方式文件 (.lnk) ---")
    lnk_paths = [
        os.path.expanduser("~/Desktop"),
        os.path.expanduser("~/Documents"),
        os.path.expanduser("~/AppData/Roaming/Microsoft/Windows/Recent"), # 最近使用的项目
    ]

    lnk_found = 0
    for path in lnk_paths:
        if not os.path.exists(path):
            continue
        for root, _, files in os.walk(path):
            for file in files:
                if file.lower().endswith(".lnk"):
                    lnk_file_path = os.path.join(root, file)
                    try:
                        # 获取LNK文件的修改时间
                        # 真正的LNK解析需要liblnk等库来提取内部的时间戳和目标路径
                        mtime = datetime.datetime.fromtimestamp(os.path.getmtime(lnk_file_path))
                        if mtime.date() == target_date.date():
                            print(f"    LNK文件: {lnk_file_path}, 最后修改时间: {format_timestamp(mtime)}")
                            lnk_found += 1
                    except Exception as e:
                        # print(f"  读取LNK文件 {lnk_file_path} 失败: {e}")
                        pass # 忽略错误
    print(f"--- 找到 {lnk_found} 个 {target_date.strftime('%Y-%m-%d')} 相关快捷方式文件 ---")


def main():
    """主程序入口"""
    print("--- Windows 日常活动取证工具 ---")

    # 检查是否有命令行参数
    if len(sys.argv) < 2:
        print("用法: python your_script_name.py YYYY-MM-DD")
        print("例如: python activity_monitor.py 2024-06-19")
        sys.exit(1) # 退出程序，表示错误

    date_str = sys.argv[1] # 获取第一个命令行参数作为日期字符串

    try:
        target_date = datetime.datetime.strptime(date_str, "%Y-%m-%d")
    except ValueError:
        print(f"错误: 日期格式不正确。请使用 YYYY-MM-DD 格式。您输入的是: {date_str}")
        sys.exit(1) # 退出程序，表示错误

    print(f"\n正在查询 {target_date.strftime('%Y年%m月%d日')} 的计算机活动...")

    # 执行各项活动查询
    get_event_logs_for_day(target_date)
    get_prefetch_files_for_day(target_date)
    get_lnk_files_for_day(target_date)

    print("\n--- 查询完成 ---")
    print("请注意：本工具仅为示例，无法获取所有活动信息。")
    print("更全面的取证分析需要专业的工具和更深入的系统知识。")

if __name__ == "__main__":
    main()